using System.Net;

/// <summary>
/// Ip限制中间件
/// </summary>
public class IpRestrictionMiddleware
{
    private readonly RequestDelegate _next;
    private readonly IConfiguration _configuration;
    private readonly ILogger<IpRestrictionMiddleware> _logger;

    public IpRestrictionMiddleware(RequestDelegate next, IConfiguration configuration, ILogger<IpRestrictionMiddleware> logger)
    {
        _next = next;
        _configuration = configuration;
        _logger = logger;
    }

    public async Task InvokeAsync(HttpContext context)
    {
        // 只对Swagger相关路径进行IP限制
        var path = context.Request.Path.Value?.ToLower();
        if (path != null && (path.Contains("/swagger") || path.Contains("/knife4j")))
        {
            var remoteIp = GetClientIpAddress(context);
            var allowedIps = _configuration.GetSection("IpRestriction:AllowedIPs").Get<string[]>() ?? new string[0];
            var allowedNetworks = _configuration.GetSection("IpRestriction:AllowedNetworks").Get<string[]>() ?? new string[0];

            if (!IsIpAllowed(remoteIp, allowedIps, allowedNetworks))
            {
                _logger.LogWarning($"IP {remoteIp} 尝试访问接口文档被拒绝");
                context.Response.StatusCode = 403;
                await context.Response.WriteAsync("Access denied from this IP address.");
                return;
            }
        }

        await _next(context);
    }

    private string GetClientIpAddress(HttpContext context)
    {
        // 处理反向代理的情况
        var xForwardedFor = context.Request.Headers["X-Forwarded-For"].FirstOrDefault();
        if (!string.IsNullOrEmpty(xForwardedFor))
        {
            return NormalizeIpAddress(xForwardedFor.Split(',')[0].Trim());
        }

        var xRealIp = context.Request.Headers["X-Real-IP"].FirstOrDefault();
        if (!string.IsNullOrEmpty(xRealIp))
        {
            return NormalizeIpAddress(xRealIp);
        }

        var remoteIp = context.Connection.RemoteIpAddress?.ToString();
        return NormalizeIpAddress(remoteIp ?? "unknown");
    }

    private string NormalizeIpAddress(string ipAddress)
    {
        if (string.IsNullOrEmpty(ipAddress)) return ipAddress;
        
        // 处理IPv6映射的IPv4地址
        if (ipAddress.StartsWith("::ffff:"))
        {
            // 提取IPv4部分
            return ipAddress.Substring(7);
        }
        
        // 处理IPv6地址的简化
        if (IPAddress.TryParse(ipAddress, out var parsedIp))
        {
            if (parsedIp.IsIPv4MappedToIPv6)
            {
                return parsedIp.MapToIPv4().ToString();
            }
        }
        
        return ipAddress;
    }

    private bool IsIpAllowed(string clientIp, string[] allowedIps, string[] allowedNetworks)
    {
        // 标准化客户端IP
        var normalizedClientIp = NormalizeIpAddress(clientIp);
        
        if (!IPAddress.TryParse(normalizedClientIp, out var clientIpAddress))
            return false;

        // 检查单个IP
        foreach (var allowedIp in allowedIps)
        {
            var normalizedAllowedIp = NormalizeIpAddress(allowedIp);
            if (IPAddress.TryParse(normalizedAllowedIp, out var allowedIpAddress) && 
                clientIpAddress.Equals(allowedIpAddress))
            {
                return true;
            }
        }

        // 检查网络段
        foreach (var network in allowedNetworks)
        {
            if (IsInNetwork(clientIpAddress, network))
            {
                return true;
            }
        }

        return false;
    }

    private bool IsInNetwork(IPAddress clientIp, string network)
    {
        try
        {
            var parts = network.Split('/');
            if (parts.Length != 2) return false;

            var networkIp = IPAddress.Parse(parts[0]);
            var prefixLength = int.Parse(parts[1]);

            var networkBytes = networkIp.GetAddressBytes();
            var clientBytes = clientIp.GetAddressBytes();

            if (networkBytes.Length != clientBytes.Length) return false;

            var bytesToCheck = prefixLength / 8;
            var bitsToCheck = prefixLength % 8;

            for (int i = 0; i < bytesToCheck; i++)
            {
                if (networkBytes[i] != clientBytes[i]) return false;
            }

            if (bitsToCheck > 0)
            {
                var mask = (byte)(0xFF << (8 - bitsToCheck));
                if ((networkBytes[bytesToCheck] & mask) != (clientBytes[bytesToCheck] & mask))
                    return false;
            }

            return true;
        }
        catch
        {
            return false;
        }
    }
}